
import copy
import json
import math
import random
import sys
import time

import torch

model = torch.jit.load("data/gomoku_230201.pth")
TIME_LIMIT = 3.8
MAX_EVAL = 500
ACTION_SIZE = 15*15

DIR = ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1))

random.seed(2023, 2)

class Game:

    board_hash = [
        [random.getrandbits(48) for _ in range(15**2)],
        [random.getrandbits(48) for _ in range(15**2)]
    ]
    init_hash = random.getrandbits(48)

    def __init__(self):
        self.board = torch.zeros((2, 15, 15), dtype=torch.bool)
        self.board_hash = Game.init_hash
        self.stones = 0

    def place(self, x, y, color):
        assert color == self.stones % 2
        self.step(x * 15 + y)

    def step(self, a):
        next_player = self.stones % 2
        self.board_hash ^= Game.board_hash[next_player][a]
        assert self.board[next_player, a//15, a%15] == False
        self.board[next_player, a//15, a%15] = True
        self.stones += 1

    def __hash__(self) -> int:
        return self.board_hash

    def to_nnet_input(self):
        next_player = self.stones % 2
        if next_player == 0:
            return self.board[None].float()
        else:
            return self.board[None, [1, 0]].float()

    def invalid_mask(self):
        return self.board.any(0)


class MCTS:
    def __init__(self, nnet):
        self.nnet = nnet

        self.Qs = {}        # stores Q values for s,a (as defined in the paper)
        self.Qsa = {}       # stores Q values for s,a (as defined in the paper)
        self.Ns = {}        # stores #times board s was visited

    def best_move(self, game, timeout=TIME_LIMIT):
        tic = time.time()
        s = hash(game)
        while time.time() - tic < timeout:
            v = self._search(copy.deepcopy(game))
            # print(v)
            if self.Ns[s] >= MAX_EVAL:
                break
        sys.stderr.write(f'Searched {self.Ns[s]} moves in {time.time() - tic} sec\n')
        return self.Qsa[s].argmax().item()

    def _search(self, game: Game):
        s = hash(game)

        if s not in self.Qsa:  # leaf node
            Qsa = self.nnet(game.to_nnet_input()).flatten()
            invalid_mask = game.invalid_mask().flatten()
            Qsa = Qsa.masked_fill_(invalid_mask, float('-inf'))
            Qs = torch.sum(torch.softmax(Qsa / 0.01, 0) * Qsa.sigmoid())
            # Qs = Qsa.max()

            self.Qsa[s] = Qsa
            self.Qs[s] = Qs
            self.Ns[s] = 1
            return Qs

        Qs = self.Qs[s]
        if Qs < 0.01 or Qs > 0.99:
            self.Ns[s] = math.inf
            return Qs
        if game.stones == 225:
            self.Ns[s] = math.inf
            return torch.tensor(0.5)

        Qsa = self.Qsa[s]
        a = torch.distributions.Categorical(logits=Qsa / 0.01).sample().item()
        game.step(a)
        v = 1 - self._search(game)
        Qsa[a] = torch.log(v.clamp_min(1e-5) / (1 - v).clamp_min(1e-5))
        Qs = torch.sum(torch.softmax(Qsa / 0.01, 0) * Qsa.sigmoid())
        self.Qs[s] = Qs
        self.Ns[s] += 1
        return Qs


def initBoard():
    fullInput = json.loads(input())
    requests = fullInput["requests"]
    responses = fullInput["responses"]
    game = Game()
    my_color = 0
    if requests[0]["x"] >= 0:
        my_color = 1
        game.place(requests[0]["x"], requests[0]["y"], 1 - my_color)
    turn = len(responses)
    for i in range(turn):
        game.place(responses[i]["x"], responses[i]["y"], my_color)
        game.place(requests[i + 1]["x"], requests[i + 1]["y"], 1 - my_color)
    return game, my_color


@torch.no_grad()
def main_plain():
    game, my_color = initBoard()

    while True:
        p = model(game.to_nnet_input())[0]
        p = p.masked_fill(game.invalid_mask(), float('-inf')).flatten()
        move = torch.distributions.Categorical(logits=p / 0.01).sample().item()
        x, y = move // 15, move % 15

        print(json.dumps({"response": {"x": x, "y": y}}))
        game.place(x, y, my_color)
        print(">>>BOTZONE_REQUEST_KEEP_RUNNING<<<", flush=True)

        request = json.loads(input())
        game.place(request["x"], request["y"], 1 - my_color)


@torch.no_grad()
def main_MCTS():
    game, my_color = initBoard()
    mcts = MCTS(model)
    while True:
        move = mcts.best_move(game)
        x, y = move // 15, move % 15

        print(json.dumps({"response": {"x": x, "y": y}}))
        game.place(x, y, my_color)
        print(">>>BOTZONE_REQUEST_KEEP_RUNNING<<<", flush=True)

        request = json.loads(input())
        game.place(request["x"], request["y"], 1 - my_color)


if __name__ == '__main__':
    main_MCTS()
